import pandas as pd
import random
import time
import direct
from maze import Maze
import pygame
import numpy as np
import os


class Environment:
    def __init__(self, name):
        # 初始化参数
        self.maze_size = tuple((10, 10))
        self.epsilon = 0.9  # 贪婪度
        self.alpha = 0.01  # 学习率
        self.gamma = 0.9  # 奖励递减值

        self.isReword = True

        #####探索者的状态,即可到达的位置
        self.actions = direct.DIRECTS

        self.readOnlyDataPath = f"./model/{name}_table.xlsx"  # 用于查看
        self.dataPath = f"./model/{name}_table.json"  # 用于存储和读取

        # q表
        self.q_table = pd.DataFrame(
            data=[[0 for _ in self.actions] for _ in range(self.maze_size[0] * self.maze_size[1])],
            index=[str(tuple((i, j))) for j in range(self.maze_size[1]) for i in range(self.maze_size[0])],
            columns=[str(dr) for dr in self.actions])

        # 生成迷宫的二维数组
        self.maze = Maze(name, self.maze_size[0], self.maze_size[1], (0, 0),
                         (self.maze_size[0] - 1, self.maze_size[1] - 1), 0.5, 3)

    def load(self):
        self.maze.load_maze()

    def save_q_table(self):
        self.q_table.to_excel(self.readOnlyDataPath)
        self.q_table.to_json(self.dataPath)

    def load_q_table(self):
        if os.path.exists(self.dataPath):
            self.q_table = pd.read_json(self.dataPath)

    def get_reward(self, state):
        t = -1
        if self.isReword:
            t = 1
        if state == self.maze.end:
            return 100
        elif state in self.maze.reward_blocks:
            return t*50
        else:
            return -1

    def get_valid_actions(self, state):
        actions = self.maze.get_block_actions(state)
        valid_actions = np.where(actions == 1)[0]
        return [str(dr) for dr in [self.actions[i] for i in valid_actions]]

    # 寻找动作的依据为以一定概率选择目前状态下动作值函数最大的动作，以一定概率随机选择（随机选择的目的是增加探索率）
    def choose_action(self, current_state):
        cur_actions = self.maze.get_block_actions(current_state)
        valid_actions = np.where(cur_actions == 1)[0]
        if random.uniform(0, 1) > self.epsilon or (self.q_table.loc[str(current_state)] == 0).all():
            return self.actions[random.choice(valid_actions)]
        else:
            columns = [str(dr) for dr in [self.actions[i] for i in valid_actions]]
            return self.actions[self.q_table.columns.get_loc(self.q_table.loc[str(current_state), columns].idxmax())]

    # 将该动作执行，得到奖励值，下个状态以及是否结束寻路标志
    def step(self, current_state, action) -> tuple:
        next_state = (current_state[0] + action[0], current_state[1] + action[1])
        reward = self.get_reward(next_state)
        done = False
        if next_state == self.maze.end:
            done = True
        return tuple((next_state, reward, done))

    # 更新q表
    def learn(self, current_state, action, reward, next_state):
        next_state_q_value = self.q_table.loc[str(next_state), self.get_valid_actions(next_state)]
        self.q_table.loc[str(current_state), str(action)] += self.alpha * (
                reward + self.gamma * next_state_q_value.max() - self.q_table.loc[str(current_state), str(action)])
        return self.q_table.loc[str(current_state), str(action)]


if __name__ == '__main__':
    # 初始化Pygame库
    pygame.init()

    # 窗口大小
    size = (500, 500)

    # 设置标题和窗口大小
    pygame.display.set_caption("Maze Game")
    #设置模型名称
    env = Environment("q1")
    print("读入模型")
    env.load_q_table()
    # 事件循环
    running = True
    #配置部分参数
    env.epsilon = 1  # 贪婪度
    # env.isReword = True  # 奖励点
    env.isReword = False  # 惩罚点
    for i in range(1000):

        env.load()

        # QLearning
        cur_state = env.maze.start

        cost = 0
        step_num = 0

        while True and running:

            for event in pygame.event.get():
                if event.type == pygame.QUIT:  # 点击关闭按钮
                    running = False
            pygame.display.flip()
            action = env.choose_action(cur_state)
            next_state, reward, done = env.step(cur_state, action)
            cost += env.learn(cur_state, action, reward, next_state)

            cur_state = next_state
            env.maze.update_bot(cur_state)
            step_num += 1
            pygame.time.wait(20)

            if done:
                print(f"第{i}次结果{step_num}")
                break

        if not running:
            print("关闭窗口")
            break

    print("已保存模型")
    env.save_q_table()

    # 退出Pygame库
    pygame.quit()
